import numpy
import Genefinder
from math import log
from numpy import loadtxt, savetxt
from Genefinder import pi, A
import parser
import math
import time
import sys

def log(num):
    if num == 0:
        return float('-inf')
    
    return math.log(num)

def viterbi(genome):
    AMatrix = loadtxt("A.txt", delimiter=',')
    pMatrix = loadtxt("p.txt", delimiter=',')
    genome = parser.Parser(genome)
    
    omega = numpy.arange(16.0*genome.length).reshape(16, genome.length)
    for e in omega.flat:
        omega.itemset(e,0)
    symbol = genome.nextsymbol()

    #Base case
    for i in range(16):
        if A[0,i]!=0:
            omega[i,0] = log(pi[i,0]) + log(pMatrix[Genefinder.mapsymbol(symbol), i])        

    #Recursion case:
    #prob = 0
    for i in range(1,genome.length):
        for j in range(16):
            prob = float("-inf")
            for z in range(16):
                tmp = log(AMatrix[j,z]) + omega[z, i-1]
                if tmp > prob:
                    prob = tmp
            omega[j,i] = prob + log(pMatrix[Genefinder.mapsymbol(symbol),j])
        symbol = genome.nextsymbol()
            
    savetxt("omega.txt", omega, '%f', ',') 
    return omega
    
def backtrace(genome, omega = None):
    fileName = genome
    genome = parser.Parser(genome)
    genome.gotoend()
    symbol = genome.peek(0)

    print "loading files"
    AMatrix = loadtxt("A.txt", delimiter=',')
    print "A loaded"
    pMatrix = loadtxt("p.txt", delimiter=',')
    print "p loaded"
    if omega is None:
        omega = loadtxt("omega.txt", delimiter=',')
    print "loading files completed"

    arg_k = 0
    max = float("-inf")
    for k in range(16):
        possibleMax = omega[k, genome.length-1]
        if possibleMax > max:
            max = possibleMax
            arg_k = k            
    resultFile = [arg_k]
    
    for n in reversed(range(genome.length-1)):
        max = float("-inf")
        arg_k = 0
        for k in range(16):
            n_1 = resultFile[genome.length-2-n]
            possibleMax = log(pMatrix[Genefinder.mapsymbol(symbol),n_1]) + omega[k,n] + log(AMatrix[n_1,k])
            if possibleMax > max:
                max = possibleMax
                arg_k = k
        resultFile.append(arg_k)
           
        symbol = genome.prevsymbol()
    
#    with open("intermediate.txt" , "w") as f:
#        f.write("".join([str(s) for s in resultFile]))

    with open("result-%s.fa" % fileName, "w") as f:
        count = 0
        for i in reversed(range(len(resultFile))):
            arg_k = resultFile[i]
            if arg_k == 0:
                f.write("N")
            else:
                f.write("C")
            count = count +1            
            if count == 60:
                count = 0
                f.write("\n")            
      
print "running the following file: "      
print sys.argv[1]
omega = viterbi(sys.argv[1])
backtrace(sys.argv[1], omega)
